// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/json_type.h>

namespace kumo::pollux::functions::prestosql {

namespace {

class JsonExtractScalarTest : public functions::test::FunctionBaseTest {
 protected:
  VectorPtr makeVector(std::optional<std::string> json, const TypePtr& type) {
    std::optional<StringView> s = json.has_value()
        ? std::make_optional(StringView(json.value()))
        : std::nullopt;
    return make_nullable_flat_vector<StringView>({s}, type);
  }

  std::optional<std::string> jsonExtractScalar(
      std::optional<std::string> json,
      std::optional<std::string> path) {
    auto jsonVector = makeVector(json, JSON());
    auto pathVector = make_nullable_flat_vector<std::string>({path});
    auto jsonResult = evaluateOnce<std::string>(
        "json_extract_scalar(c0, c1)", make_row_vector({jsonVector, pathVector}));

    auto varcharVector = makeVector(json, VARCHAR());
    auto varcharResult = evaluateOnce<std::string>(
        "json_extract_scalar(c0, c1)",
        make_row_vector({varcharVector, pathVector}));

    EXPECT_EQ(jsonResult, varcharResult);
    return jsonResult;
  }

  void assertResults(
      const std::vector<std::optional<StringView>>& json,
      const std::vector<std::optional<StringView>>& path,
      const std::vector<std::optional<StringView>>& expected) {
    auto jsonVector = make_nullable_flat_vector<StringView>(json, JSON());
    auto pathVector = make_nullable_flat_vector<StringView>(path);
    auto expectedVector = make_nullable_flat_vector<StringView>(expected);

    pollux::test::assertEqualVectors(
        expectedVector,
        evaluate<SimpleVector<StringView>>(
            "json_extract_scalar(c0, c1)",
            make_row_vector({jsonVector, pathVector})));

    auto varcharVector = make_nullable_flat_vector<StringView>(json, VARCHAR());

    pollux::test::assertEqualVectors(
        expectedVector,
        evaluate<SimpleVector<StringView>>(
            "json_extract_scalar(c0, c1)",
            make_row_vector({varcharVector, pathVector})));
  }
};

TEST_F(JsonExtractScalarTest, simple) {
  // Scalars.
  EXPECT_EQ(jsonExtractScalar(R"(1)", "$"), "1");
  EXPECT_EQ(jsonExtractScalar(R"(123456)", "$"), "123456");
  EXPECT_EQ(jsonExtractScalar(R"("hello")", "$"), "hello");
  EXPECT_EQ(jsonExtractScalar(R"(1.1)", "$"), "1.1");
  EXPECT_EQ(jsonExtractScalar(R"("")", "$"), "");
  EXPECT_EQ(jsonExtractScalar(R"(true)", "$"), "true");

  // Simple lists.
  EXPECT_EQ(jsonExtractScalar(R"([1,2])", "$[0]"), "1");
  EXPECT_EQ(jsonExtractScalar(R"([1,2])", "$[1]"), "2");
  EXPECT_EQ(jsonExtractScalar(R"([1,2])", "$[2]"), std::nullopt);
  EXPECT_EQ(jsonExtractScalar(R"([1,2])", "$[999]"), std::nullopt);

  // Simple maps.
  EXPECT_EQ(jsonExtractScalar(R"({"k1":"v1"})", "$.k1"), "v1");
  EXPECT_EQ(jsonExtractScalar(R"({"k1":"v1"})", "$.k2"), std::nullopt);
  EXPECT_EQ(jsonExtractScalar(R"({"k1":"v1"})", "$.k1.k3"), std::nullopt);
  EXPECT_EQ(jsonExtractScalar(R"({"k1":[0,1,2]})", "$.k1"), std::nullopt);
  EXPECT_EQ(jsonExtractScalar(R"({"k1":""})", "$.k1"), "");

  // Nested
  EXPECT_EQ(jsonExtractScalar(R"({"k1":{"k2": 999}})", "$.k1.k2"), "999");
  EXPECT_EQ(jsonExtractScalar(R"({"k1":[1,2,3]})", "$.k1[0]"), "1");
  EXPECT_EQ(jsonExtractScalar(R"({"k1":[1,2,3]})", "$.k1[2]"), "3");
  EXPECT_EQ(jsonExtractScalar(R"([{"k1":"v1"}, 2])", "$[0].k1"), "v1");
  EXPECT_EQ(jsonExtractScalar(R"([{"k1":"v1"}, 2])", "$[1]"), "2");
  EXPECT_EQ(
      jsonExtractScalar(
          R"([{"k1":[{"k2": ["v1", "v2"]}]}])", "$[0].k1[0].k2[1]"),
      "v2");

  // Paths without leading '$'.
  EXPECT_EQ(jsonExtractScalar(R"({"k1":"v1"})", "k1"), "v1");
  EXPECT_EQ(jsonExtractScalar(R"({"k1":{"k2": 999}})", "k1.k2"), "999");
  EXPECT_EQ(jsonExtractScalar(R"([1,2])", "[0]"), "1");
  EXPECT_EQ(jsonExtractScalar(R"({"k1":"v1"})", R"(["k1"])"), "v1");

  // Paths with redundant '.'s.
  auto json = "[1, 2, [10, 20, [100, 200, 300]]]";
  EXPECT_EQ(jsonExtractScalar(json, "$[2][2][2]"), "300");
  EXPECT_EQ(jsonExtractScalar(json, "$.[2].[2].[2]"), "300");
  EXPECT_EQ(jsonExtractScalar(json, "$[2].[2].[2]"), "300");
  EXPECT_EQ(jsonExtractScalar(json, "$[2][2].[2]"), "300");
}

TEST_F(JsonExtractScalarTest, jsonType) {
  // Scalars.
  assertResults(
      {R"(1)"_sv,
       R"(123456)"_sv,
       R"("hello")"_sv,
       R"(1.1)"_sv,
       R"("")"_sv,
       R"(true)"},
      {"$"_sv, "$"_sv, "$"_sv, "$"_sv, "$"_sv, "$"_sv},
      {"1"_sv, "123456"_sv, "hello"_sv, "1.1"_sv, ""_sv, "true"_sv});

  // Simple lists.
  assertResults(
      {R"([1,2])"_sv, R"([1,2])"_sv, R"([1,2])"_sv, R"([1,2])"_sv},
      {"$[0]"_sv, "$[1]"_sv, "$[2]"_sv, "$[999]"_sv},
      {"1"_sv, "2"_sv, std::nullopt, std::nullopt});

  // Simple maps.
  assertResults(
      {R"({"k1":"v1"})"_sv,
       R"({"k1":"v1"})"_sv,
       R"({"k1":"v1"})"_sv,
       R"({"k1":[0,1,2]})"_sv,
       R"({"k1":""})"_sv},
      {"$.k1"_sv, "$.k2"_sv, "$.k1.k3"_sv, "$.k1"_sv, "$.k1"_sv},
      {"v1"_sv, std::nullopt, std::nullopt, std::nullopt, ""_sv});

  // Nested
  assertResults(
      {R"({"k1":{"k2": 999}})"_sv,
       R"({"k1":[1,2,3]})"_sv,
       R"({"k1":[1,2,3]})"_sv,
       R"([{"k1":"v1"}, 2])"_sv,
       R"([{"k1":"v1"}, 2])"_sv,
       R"([{"k1":[{"k2": ["v1", "v2"]}]}])"_sv},
      {"$.k1.k2"_sv,
       "$.k1[0]"_sv,
       "$.k1[2]"_sv,
       "$[0].k1"_sv,
       "$[1]"_sv,
       "$[0].k1[0].k2[1]"_sv},
      {"999"_sv, "1"_sv, "3"_sv, "v1"_sv, "2"_sv, "v2"_sv});
}

TEST_F(JsonExtractScalarTest, utf8) {
  EXPECT_EQ(
      jsonExtractScalar(R"({"k1":"I \u2665 UTF-8"})", "$.k1"),
      "I \u2665 UTF-8");
  EXPECT_EQ(
      jsonExtractScalar("{\"k1\":\"I \u2665 UTF-8\"}", "$.k1"),
      "I \u2665 UTF-8");

  EXPECT_EQ(
      jsonExtractScalar("{\"k1\":\"I \U0001D11E playing in G-clef\"}", "$.k1"),
      "I \U0001D11E playing in G-clef");
}

TEST_F(JsonExtractScalarTest, invalidPath) {
  POLLUX_ASSERT_THROW(jsonExtractScalar(R"([0,1,2])", ""), "Invalid JSON path");
  POLLUX_ASSERT_THROW(
      jsonExtractScalar(R"([0,1,2])", "$[]"), "Invalid JSON path");
  POLLUX_ASSERT_THROW(
      jsonExtractScalar(R"([0,1,2])", "$-1"), "Invalid JSON path");
  POLLUX_ASSERT_THROW(
      jsonExtractScalar(R"({"k1":"v1"})", "$k1"), "Invalid JSON path");
  POLLUX_ASSERT_THROW(
      jsonExtractScalar(R"({"k1":"v1"})", "$.k1."), "Invalid JSON path");
  POLLUX_ASSERT_THROW(
      jsonExtractScalar(R"({"k1":"v1"})", "$.k1["), "Invalid JSON path");
  POLLUX_ASSERT_THROW(
      jsonExtractScalar(R"({"k1":"v1)", "$.k1["), "Invalid JSON path");
}

TEST_F(JsonExtractScalarTest, invalidJson) {
  // Verify that we return null on invalid JSON regardless of whether the
  // invalid section is after the target path.
  EXPECT_EQ(jsonExtractScalar(R"({"a": "b", "c": "d})", "$.a"), std::nullopt);
  EXPECT_EQ(jsonExtractScalar(R"([["a"], ["b]])", "$[0][0]"), std::nullopt);
}

TEST_F(JsonExtractScalarTest, escapedString) {
  // Verify the the returned string is unescaped.
  EXPECT_EQ(
      jsonExtractScalar(R"({"x": {"a" : 1, "b" : "b\/c"} })", "$.x.b"), "b/c");
}

// simdjson, like Presto java, returns the large number as-is as a string,
// without trying to convert it to an integer.
TEST_F(JsonExtractScalarTest, overflow) {
  EXPECT_EQ(
      jsonExtractScalar(
          R"(184467440737095516151844674407370955161518446744073709551615)",
          "$"),
      "184467440737095516151844674407370955161518446744073709551615");
}

// TODO: When there is a wildcard in the json path, Presto's behavior is to
// always extract an array of selected items, and hence json_extract_scalar()
// always return NULL in this situation. But some internal customers are
// relying on the current behavior of returning the selected item itself if
// there is exactly one selected. We'll fix json_extract_scalar() to follow
// Presto's behavior after internal customers clear the dependence on this
// diverged semantic. This unit test makes sure we don't break their workloads
// before they clear the dependency.
TEST_F(JsonExtractScalarTest, wildcardSelect) {
  EXPECT_EQ(
      jsonExtractScalar(R"({"tags":{"a":["b"],"c":["d"]}})", "$.tags.c[*]"),
      "d");
  EXPECT_EQ(
      jsonExtractScalar(R"({"tags":{"a":["b"],"c":["d"]}})", "$[tags][c][*]"),
      "d");
  EXPECT_EQ(
      jsonExtractScalar(R"({"tags":{"a":["b"],"c":["d","e"]}})", "$.tags.c[*]"),
      std::nullopt);
  EXPECT_EQ(
      jsonExtractScalar(R"({"tags":{"a":["b"],"c":[]}})", "$.tags.c[*]"),
      std::nullopt);
}

} // namespace

} // namespace kumo::pollux::functions::prestosql
